package org.neu.ccs.secureim;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.DataInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.OutputStream;
import java.net.ConnectException;
import java.net.InetAddress;
import java.net.Socket;
import java.security.Key;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.spec.X509EncodedKeySpec;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Vector;
import java.util.concurrent.ConcurrentHashMap;

import javax.crypto.Cipher;
import javax.crypto.Mac;
import javax.crypto.spec.DHParameterSpec;

import org.apache.commons.codec.binary.Base64;

import com.sun.xml.internal.bind.v2.runtime.unmarshaller.XsiNilLoader.Single;


/*
 * Class to encapsulate the client and its actions.
 */

public class ChatClient {
	KeyPair asymKeys;
	public SecureRandom r;
	String user;
	char[] pass;
	InetAddress serverAddr;
	int port;
	PublicKey serverPK;
	public KeyPair clientPair;
	int clientPort;
	
	public ConcurrentHashMap<String, Key> sharedKeys;
	public ConcurrentHashMap<String, UserInfo> userLocs;
	public ConcurrentHashMap<String, String> userHashes;
	
	ClientListenerThread listenerThread;
	
	static final int LOGIN = 1;
	static final int REGISTER = 2;
	static final int OKAY = 3;
	static final int ERROR = 4;
	static final int LIST = 5;
	static final int LOCATE = 6;
	static final int CONNECT = 7;
	static final int MESSAGE = 8;
	static final int LOGOUT = 9;
	
	public ChatClient(String user, char[] pass, InetAddress serverAddr, int port, byte[] serverPK)
	{
		r = new SecureRandom();
		this.user = user;
		this.pass = pass;
		this.serverAddr = serverAddr;
		this.port = port;
		try
		{
			KeyFactory kf = KeyFactory.getInstance("RSA");
			this.serverPK = kf.generatePublic(new X509EncodedKeySpec(serverPK));
		}
		catch(Exception e)
		{
			e.printStackTrace();
			System.exit(-1);
		}
		sharedKeys = new ConcurrentHashMap<String, Key>();
		userLocs = new ConcurrentHashMap<String, UserInfo>();
		userHashes = new ConcurrentHashMap<String, String>();
	}
	
	public int start() throws IOException
	{
		listenerThread = new ClientListenerThread(this);
		int newPort =  listenerThread.initSocket();
		clientPort = newPort;
		Thread t = new Thread(listenerThread);
		t.start();
		return newPort;
	}
	
	public void register()
	{
		byte[] salt = new byte[128];
		r.nextBytes(salt);
		
		clientPair = CryptoHelper.GenerateClientKeyPair(user, pass, salt);
		byte[] publicKey = clientPair.getPublic().getEncoded();
		//TODO: Send key, username and salt to server
		//In practice, encoded public key is 294 bytes
		//Key 2048 bits, salt 128 bytes, username string
		//Message: {REGISTER <key> <salt> <username>}Spub

		try
		{
			//PackRSA takes the public key to pack with, a random number generate and then any number of byte arrays to pack, i.e. in this
			//case publicKey salt and the username are being packed with the public key serverPK
			byte[] encrypted = CryptoHelper.PackRSA(serverPK, r, publicKey, salt, user.getBytes());
			
			Socket sock = new Socket(serverAddr, port);
			if( sock.isConnected() )
			{
				OutputStream sout = sock.getOutputStream();
				sout.write(REGISTER);
				NetworkHelper.WriteInt(sout,encrypted.length);
				sout.write(encrypted);
			}
			
			sock.close();
			
		}
		catch(Exception e)
		{
			e.printStackTrace();
		} 
	}
	
	public void testConnection()
	{
		try {
			Socket sock = new Socket(serverAddr, port);
			if( sock.isConnected() )
				System.out.println("Connection successfull!");
		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
	}
	
	public boolean login()
	{
		//TODO: Request salt from server
		//Challenge 128 bytes timestamp System.nanoTime() 64-bit number
		//Message: {LOGIN <username> <R1>}Spub
		//Return: {<salt> <R1>}Spriv {<R2> <T1> <username>}Spriv
		//Message2: {R2}Apriv {<R2> <T1> <username>}Spriv
			//byte[] salt = null;
			//clientPair = CryptoHelper.GenerateClientKeyPair(user, pass, salt);
		//TODO: Login protocol 
		//Key 2048 bits, salt 128 bytes, username
		
		try
		{
			Socket sock = new Socket(serverAddr, port);
			if( sock.isConnected() )
			{
				OutputStream sout = sock.getOutputStream();				
				
				byte[] rand1 = new byte[128];
				r.nextBytes(rand1);
				
				byte[] messageOne = CryptoHelper.PackRSA(serverPK,r,user.getBytes(), rand1);
				sout.write(LOGIN);
				NetworkHelper.WriteInt(sout,messageOne.length);
				sout.write(messageOne);
				// RECEIVE MESSAGE FROM SERVER
				
				DataInputStream sin = new DataInputStream(sock.getInputStream());
				
				int length = NetworkHelper.ReadInt(sin);
				
				if( length > 1500 )
				{
					System.out.println("Unexpected response");
					return false;
				}
				
				byte[] received = new byte[length];
				sin.readFully(received);
				// salt + R1  + R2  + T1 + user.namelength
				// 128  + 128 + 128 + 64 + user.nameLength
				
				byte[] part1 = Arrays.copyOfRange(received, 0, 512);
				byte[] part2 = Arrays.copyOfRange(received, 512, length);
				
				// Verify rand1, if OK continue, else return
				
				if( !(CryptoHelper.VerifyRSA(serverPK, part1) && CryptoHelper.VerifyRSA(serverPK, part2) && Arrays.equals(rand1, Arrays.copyOfRange(part1, 384, 512))) )
				{
					System.out.println("Login failed, incorrect key");
					sout.write(ERROR);
					return false;
				}
				
				// pull salt out
				byte[] salt = new byte[128];
				salt = Arrays.copyOfRange(part1, 256, 384);
				
				// pull rand2 out
				byte[] rand2 = new byte[128];
				rand2 = Arrays.copyOfRange(part2, 256, 384);
				
				// pull cookie out
				
				clientPair = CryptoHelper.GenerateClientKeyPair(user, pass, salt);
				
				// Send rand2 to server encrypted with Apriv
				System.out.println("Logging in on port " + clientPort);
				byte[] messageTwo = CryptoHelper.SignRSA(clientPair.getPrivate(), NetworkHelper.IntToBytes(clientPort), rand2);
				NetworkHelper.WriteInt(sout,messageTwo.length + part2.length);
				sout.write(messageTwo);
				sout.write(part2);
				
				int res = sin.read();
				if( res == OKAY )
				{
					System.out.println("Logged in successfully in!");
					return true;
				}
				else
				{
					System.out.println("Login failed.");
					return false;
				}
			}
			return false;
		}
		catch(Exception e)
		{
			e.printStackTrace();
			return false;
		}
		
		
		
	}
	
	@SuppressWarnings("unchecked")
	public Vector<String> list()
	{
		//TODO: Get list of users online from server and return it
		//Message: {LIST, T1}Apriv
		//Return: {String[] Users} Signed with Spriv
		try
		{
			Socket sock = new Socket(serverAddr, port);
			if( sock.isConnected() )
			{
				OutputStream sout = sock.getOutputStream();
				ByteArrayOutputStream out = new ByteArrayOutputStream();
				NetworkHelper.WriteInt(out, (int)System.currentTimeMillis());
				byte[] signed = CryptoHelper.SignRSA(clientPair.getPrivate(), out.toByteArray());
				sout.write(LIST);
				
				byte[] usernamePacked = CryptoHelper.PackRSA(serverPK, r, user.getBytes());
				NetworkHelper.WriteInt(sout, usernamePacked.length);
				sout.write(usernamePacked);
				
				NetworkHelper.WriteInt(sout, signed.length);
				sout.write(signed);
				
				DataInputStream sin = new DataInputStream(sock.getInputStream());
				int response = sin.read();
				if( response == ERROR )
				{
					System.out.println("Request failed (possibly user not online)");
					return null;
				}
				int length = NetworkHelper.ReadInt(sin);
				
				if( length > 1500 )
				{
					System.out.println("Unexpected response");
					return null;
				}
				
				byte[] list = new byte[length];
				sin.readFully(list);
				byte[] encrypted = Arrays.copyOfRange(list, 0, length);
				ObjectInputStream deserializer = new ObjectInputStream(new ByteArrayInputStream(CryptoHelper.UnpackRSA(clientPair.getPrivate(), encrypted)));
				return (Vector<String>)deserializer.readObject();
			}
			return null;
		}
		catch(Exception e)
		{
			e.printStackTrace();
			return null;
		}
	}
	
	public UserInfo locate(String userName)
	{
		//TODO: Get user location and public key from server
		//Message: {LOCATE <username>}Apriv
		//Return: {<public key><ip address>} Signed with Apub
		try
		{
			Socket sock = new Socket(serverAddr, port);
			if( sock.isConnected() )
			{
				OutputStream sout = sock.getOutputStream();
				ByteArrayOutputStream out = new ByteArrayOutputStream();
				out.write(userName.getBytes());
				byte[] signed = CryptoHelper.SignRSA(clientPair.getPrivate(), out.toByteArray());
				sout.write(LOCATE);
				
				byte[] usernamePacked = CryptoHelper.PackRSA(serverPK, r, user.getBytes());
				NetworkHelper.WriteInt(sout, usernamePacked.length);
				sout.write(usernamePacked);
				
				NetworkHelper.WriteInt(sout, signed.length);
				sout.write(signed);
				
				DataInputStream sin = new DataInputStream(sock.getInputStream());
				int response = sin.read();
				if( response == ERROR )
				{
					System.out.println("Request failed");
					return null;
				}
				int length = NetworkHelper.ReadInt(sin);
				
				if( length > 1500 )
				{
					System.out.println("Unexpected response");
					return null;
				}
				
				byte[] list = new byte[length];
				sin.readFully(list);
				byte[] encrypted = Arrays.copyOfRange(list, 0, length);
				byte[] unencrypted = CryptoHelper.UnpackRSA(clientPair.getPrivate(), encrypted);
				ObjectInputStream deserializer = new ObjectInputStream(new ByteArrayInputStream(unencrypted));
				
				InetAddress address = (InetAddress)deserializer.readObject();
				int newPort = deserializer.readInt();
				PublicKey key = (PublicKey)deserializer.readObject();
				
				return new UserInfo(address, newPort, key, userName);
			}
			return null;
		}
		catch(Exception e)
		{
			e.printStackTrace();
			return null;
		}
	}
	
	public Key connect(String user)
	{
		return connect(locate(user));
	}
	
	public Key connect(UserInfo info)
	{
		try
		{
			Socket sock = new Socket(info.getLocation(), info.getPort());
			if( sock.isConnected() )
			{
				DHValues vals = new DHValues(r);
				OutputStream sout = sock.getOutputStream();
				sout.write(CONNECT);
				
				byte[] userBytes = user.getBytes();
				byte[] encryptedUsername = CryptoHelper.PackRSA(info.getPubKey(), r, userBytes);
				NetworkHelper.WriteInt(sout, encryptedUsername.length);
				sout.write(encryptedUsername);
				
				ByteArrayOutputStream out = new ByteArrayOutputStream();
				ObjectOutputStream serializer = new ObjectOutputStream(out);
				
				DHParameterSpec spec = vals.getParameterSpec();
				serializer.writeObject(spec.getP());
				serializer.writeObject(spec.getG());
				serializer.writeInt(spec.getL());
				serializer.writeObject(vals.getPublicKey());
				serializer.flush();
				
				byte[] message = out.toByteArray();
				byte[] signed = CryptoHelper.SignRSA(clientPair.getPrivate(), message);
				NetworkHelper.WriteInt(sout, signed.length);
				sout.write(signed);
				
				DataInputStream sin = new DataInputStream(sock.getInputStream());
				int length = NetworkHelper.ReadInt(sin);
				
				if( length > 1500 )
				{
					System.out.println("Unexpected response");
					return null;
				}
				
				byte[] message2 = new byte[length];
				sin.readFully(message2);
				
				if( !CryptoHelper.VerifyRSA(info.getPubKey(), message2) )
				{
					System.out.println("Signature verification failed.");
					return null;
				}
				
				ByteArrayInputStream in = new ByteArrayInputStream(Arrays.copyOfRange(message2, 256, message2.length));
				ObjectInputStream deserializer = new ObjectInputStream(in);
				PublicKey pk = (PublicKey)deserializer.readObject();
				Key symKey = vals.doFinal(pk, "AES");
				
				return symKey;
			}
			else
			{
				System.out.println("Could not connect to client");
			}
			return null;
		}
		catch(ConnectException e)
		{
			System.out.println("Target client is not online");
			return null;
		}
		catch(Exception e)
		{
			e.printStackTrace();
			return null;
		}
	}
	
	public void send(String userName, String message) throws Exception
	{
		UserInfo info = null;
		if( userLocs.containsKey(userName) )
		{
			info = userLocs.get(userName);
		}
		else
		{
			info = locate(userName);
			if( info == null )
				return;
			userLocs.put(userName, info);
		}
		
		Key symKey = null;
		if( sharedKeys.containsKey(userName) )
			symKey = sharedKeys.get(userName);
		else
		{
			symKey = connect(info);
			if( symKey == null )
				return;
			sharedKeys.put(userName, symKey);
			Mac mac = Mac.getInstance("HmacSHA1");
			mac.init(symKey);
			mac.update(userName.getBytes());
			byte[] hash = mac.doFinal();
			userHashes.put(Base64.encodeBase64String(hash), userName);
		}
		
		if( symKey == null || info == null )
			return;
		
		try
		{
			Socket sock = new Socket(info.getLocation(), info.getPort());
			
			if(sock.isConnected())
			{
				OutputStream sout = sock.getOutputStream();				
				sout.write(MESSAGE);
				
				// ENCRYPT MESSAGE WITH SHARED SESSION KEY
				Cipher cipher = Cipher.getInstance("AES");
				cipher.init(Cipher.ENCRYPT_MODE, symKey);
	            byte[] encrypted = cipher.doFinal(message.getBytes());
				
				
				// Generate HMAC(KEY, MESSAGE)
				Mac m = Mac.getInstance("HmacSHA1");
				m.init(symKey);
				m.update(message.getBytes());
				byte[] mac = m.doFinal();
				
				m.update(user.getBytes());
				byte[] userHash = m.doFinal();
				
				NetworkHelper.WriteInt(sout, userHash.length);
				sout.write(userHash);
				
				// Combine lengths and send
				// Send message
				// Send HMAC
				NetworkHelper.WriteInt(sout,encrypted.length);
				sout.write(encrypted);
				NetworkHelper.WriteInt(sout,mac.length);
				sout.write(mac);
			}
			return; 
		}
		catch(ConnectException e)
		{
			System.out.println("Target client is not online");
			return;
		}
		catch(Exception e)
		{
			e.printStackTrace();
			return;
		}
	}
	
	public void logout()
	{
		try
		{
			Socket sock = new Socket(serverAddr, port);
			
			if(sock.isConnected())
			{
				OutputStream sout = sock.getOutputStream();	
				InputStream sin = sock.getInputStream();
				sout.write(LOGOUT);
				
				byte[] usernamePacked = CryptoHelper.PackRSA(serverPK, r, user.getBytes());
				NetworkHelper.WriteInt(sout, usernamePacked.length);
				sout.write(usernamePacked);
				
				byte[] signed = CryptoHelper.SignRSA(clientPair.getPrivate(), NetworkHelper.IntToBytes((int)System.currentTimeMillis()));
				NetworkHelper.WriteInt(sout, signed.length);
				sout.write(signed);
				
				if( sin.read() == OKAY )
					System.out.println("Logged out");
				else
					System.out.println("Error logging out");
			}
		}
		catch(Exception e)
		{
			
		}
	}
}
